package org.apache.flink.runtime.entrypoint.master;

import java.util.ArrayList;
import java.util.Collection;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Supplier;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;

import org.apache.flink.runtime.clusterframework.ApplicationStatus;
import org.apache.flink.runtime.leaderretrieval.LeaderRetrievalService;
import org.apache.flink.runtime.resourceconsumer.runner.ResourceConsumerRunner;
import org.apache.flink.runtime.resourcemanager.ResourceManager;
import org.apache.flink.runtime.resourcemanager.ResourceManagerService;
import org.apache.flink.runtime.rest.RestService;
import org.apache.flink.runtime.rpc.FatalErrorHandler;
import org.apache.flink.util.AutoCloseableAsync;
import org.apache.flink.util.ExceptionUtils;
import org.apache.flink.util.ERException;
import org.apache.flink.util.concurrent.FutureUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * @类名: MasterComponent
 * @说明: master组件持有者
 *       包括集群组件和rest组件，用户扩展，支持自定义组件
 *       持有/关闭清理
 *
 * @author   leehom
 * @Date	 2023年12月18日 下午7:08:30
 * 修改记录：
 *   > 增加资源消费者
 *
 * @see 	 
 */
public class MasterComponentHolder implements MasterComponentOperations {

    private static final Logger LOG =
            LoggerFactory.getLogger(MasterComponentHolder.class);

    /*********************************/
    /** 资源消费者运行管理服务*/
    @Nonnull private final ResourceConsumerRunner resourceConsumerRunner;
    /** 资源消费者主节点获取服务*/
    @Nonnull private final LeaderRetrievalService resourceConsumerRetrievalService;
    /** 资源管理器运行管理服务*/
    @Nonnull private final ResourceManagerService resourceManagerService;
    /** 资源管理器主节点获取服务*/
    @Nonnull private final LeaderRetrievalService resourceManagerRetrievalService;
    
    /** rest服务*/
    @Nonnull private final RestService webMonitorEndpoint;

    /** 终止和关闭*/
    private final CompletableFuture<Void> terminationFuture;
    private final CompletableFuture<ApplicationStatus> shutDownFuture;

    private final AtomicBoolean isRunning = new AtomicBoolean(true);

    private final FatalErrorHandler fatalErrorHandler;

    public MasterComponentHolder(
    		@Nonnull ResourceConsumerRunner resourceConsumerRunner,
            @Nonnull LeaderRetrievalService resourceConsumerRetrievalService,
            @Nonnull ResourceManagerService resourceManagerService,
            @Nonnull LeaderRetrievalService resourceManagerRetrievalService,
            @Nonnull RestService webMonitorEndpoint,
            @Nonnull FatalErrorHandler fatalErrorHandler) {
    	this.resourceConsumerRunner = resourceConsumerRunner;
    	this.resourceConsumerRetrievalService = resourceConsumerRetrievalService;
        this.resourceManagerService = resourceManagerService;
        this.resourceManagerRetrievalService = resourceManagerRetrievalService;
        this.webMonitorEndpoint = webMonitorEndpoint;
        this.fatalErrorHandler = fatalErrorHandler;
        //
        this.terminationFuture = new CompletableFuture<>();
        this.shutDownFuture = new CompletableFuture<>();

        registerShutDownFuture();
        handleUnexpectedResourceManagerTermination();
    }

    private void handleUnexpectedResourceManagerTermination() {
        resourceManagerService
                .getTerminationFuture()
                .whenComplete(
                        (ignored, throwable) -> {
                            if (isRunning.get()) {
                                fatalErrorHandler.onFatalError(
                                        new ERException(
                                                "Unexpected termination of ResourceManagerService.",
                                                throwable));
                            }
                        });
    }

    private void registerShutDownFuture() {
        // FutureUtils.forward(resourceConsumerRunner.getResultFuture(), shutDownFuture);
    }

    public final CompletableFuture<ApplicationStatus> getShutDownFuture() {
        return shutDownFuture;
    }

    /**
     * Deregister the application from the resource management system by signalling the {@link
     * ResourceManager} and also stop the process.
     *
     * @param applicationStatus to terminate the application with
     * @param diagnostics additional information about the shut down, can be {@code null}
     * @return Future which is completed once the shut down
     */
    public CompletableFuture<Void> stopApplication(
            final ApplicationStatus applicationStatus, final @Nullable String diagnostics) {
        return internalShutdown(
                () -> resourceManagerService.deregisterApplication(applicationStatus, diagnostics));
    }

    /**
     * Close the cluster components. This method will not deregister the 
     * application from the resource management and only stop the process.
     *
     * @return Future which is completed once the shut down
     */
    public CompletableFuture<Void> stopProcess() {
        return internalShutdown(FutureUtils::completedVoidFuture);
    }

    // 关闭组件和服务
    private CompletableFuture<Void> internalShutdown(
            final Supplier<CompletableFuture<?>> additionalShutdownAction) {
        if (isRunning.compareAndSet(true, false)) {
        	// 
        	final CompletableFuture<Void> webMonitorShutdownFuture =
        					webMonitorEndpoint.closeAsync();
        	// 
            final CompletableFuture<Void> closeWebMonitorAndAdditionalShutdownActionFuture =
                    FutureUtils.composeAfterwards(
                            webMonitorShutdownFuture, additionalShutdownAction);
            // 关闭服务
            return FutureUtils.composeAfterwards(
                    closeWebMonitorAndAdditionalShutdownActionFuture, this::closeAsyncInternal);
        } else {
            return terminationFuture;
        }
    }

    // 异步close，内部方法
    private CompletableFuture<Void> closeAsyncInternal() {
        LOG.info("Closing components.");

        Exception exception = null;
        // 组件关闭futures集合
        final Collection<CompletableFuture<Void>> terminationFutures = new ArrayList<>(3);
        // 停止主节点获取服务
        try {
            resourceConsumerRetrievalService.stop();
        } catch (Exception e) {
            exception = ExceptionUtils.firstOrSuppressed(e, exception);
        }
        try {
            resourceManagerRetrievalService.stop();
        } catch (Exception e) {
            exception = ExceptionUtils.firstOrSuppressed(e, exception);
        }
        //
        terminationFutures.add(resourceConsumerRunner.closeAsync());
        terminationFutures.add(resourceManagerService.closeAsync());
        if (exception != null) {
            terminationFutures.add(FutureUtils.completedExceptionally(exception));
        }
        // 
        final CompletableFuture<Void> componentTerminationFuture =
                							FutureUtils.completeAll(terminationFutures);
        //
        componentTerminationFuture.whenComplete(
                (aVoid, throwable) -> {
                    if (throwable != null) {
                        terminationFuture.completeExceptionally(throwable);
                    } else {
                        terminationFuture.complete(aVoid);
                    }
                });

        return terminationFuture;
    }

    // 异步关闭master组件
    @Override
    public CompletableFuture<Void> closeAsync() {
        return stopApplication(
                ApplicationStatus.CANCELED, "MasterComponent has been closed.");
    }
    
    public int getRestPort() {
        return webMonitorEndpoint.getRestPort();
    }

}
